from spinup import ppo_pytorch as ppo
from spinup import ddpg_pytorch as ddpg
from spinup import maddpg_pytorch as maddpg
# from spinup.utils.test_policy import load_policy_and_env, run_policy
from test_maddpg  import load_policy_and_env, run_policy
import torch
import gym

TRAIN = 1

env = lambda : gym.make('ur5controllerdual-v2')


if TRAIN:
    ac_kwargs = dict(hidden_sizes=[128,128,128], activation=torch.nn.ReLU)
    logger_kwargs = dict(output_dir='logDualUR5', exp_name='DualUR5')

    maddpg(env, ac_kwargs=ac_kwargs, logger_kwargs=logger_kwargs,
           steps_per_epoch=200, epochs=1000, replay_size=int(1e6), gamma=0.95,
           polyak=0.995, pi_lr=1e-3, q_lr=1e-3, batch_size=1024, start_steps=1900000,
           update_after=100, update_every=50, act_noise=0.1, num_test_episodes=10,
           max_ep_len=100, save_freq=1)

else:
    _, get_action1, get_action2 = load_policy_and_env('logDualUR5')
    env_test = gym.make('ur5controllerdual-v2')
    run_policy(env_test, get_action1, get_action2)
